import json
import base64
import datetime
import hmac
import copy
import random
import time
from dataclasses import dataclass, field
from enum import Enum
from hashlib import sha256
from typing import Optional, List, Any, Dict

# picc中台请求报文格式
class SessionStatus(Enum):
    ONCE = "once"
    BEGIN = "begin"
    CONTINUE = "continue"
    END = "end"

    @classmethod
    def value_of(cls, value: str):
        for mode in cls:
            if mode.value == value:
                return mode
        raise ValueError(f"invalid session status type value {value}")

class MessageRole(Enum):
    USER = "user"
    SYSTEM = "system"
    ASSISTANT = "assistant"

    @classmethod
    def value_of(cls, value: str):
        for mode in cls:
            if mode.value == value:
                return mode
        raise ValueError(f"invalid messsage role type value {value}")

class PiccModel(Enum):
    QWEN = "qwen"
    DEEPSEEK = "deepseek"

    @classmethod
    def value_of(cls, value: str):
        """
        Get value of given mode.

        :param value: mode value
        :return: mode
        """
        for mode in cls:
            if mode.value == value:
                return mode
        raise ValueError(f"invalid Picc Model type value {value}")

@dataclass
class Session:
    sid: str = ""
    following: str = ""
    status: SessionStatus = SessionStatus.ONCE

    def to_dict(self):
        return {
            "sid": self.sid,
            "following": self.following,
            "status": self.status.value
        }

@dataclass
class ProductParam:
    product_id: str = ""
    scene_id: str = ""
    user_id: str = ""
    request_id: str = ""
    extend_product_params: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self):
        return {
            "product_id": self.product_id,
            "scene_id": self.scene_id,
            "user_id": self.user_id,
            "request_id": self.request_id,
            "extend_product_params": self.extend_product_params
        }

@dataclass
class ModelParams:
    pass

@dataclass
class Params:
    product_params: ProductParam
    model_params: ModelParams

    def to_dict(self):
        return {
            "product_params": self.product_params.to_dict(),
            "model_params": self.model_params.to_dict()
        }

@dataclass
class Message:
    role: MessageRole
    content: str

    def to_dict(self):
        return {
            "role": self.role.value,
            "content": self.content
        }

@dataclass
class Messages:
    text: list[Message] = field(default_factory=list)

    def to_dict(self):
        return {
            "text": [msg.to_dict() for msg in self.text[-7:]],
        }
@dataclass
class Payloads:
    message: list[Message] = field(default_factory=list)

    def to_dict(self):
        return {
            "message": [msg.to_dict() for msg in self.message],
        }

@dataclass
class PayloadsWithTools(Payloads):
    tools: list[str] = field(default_factory=list)

    def to_dict(self):
        dic = super().to_dict()
        dic["tools"] = [tool for tool in self.tools]
        return dic

@dataclass
class PiccRequest:
    session: Session = field(default_factory=Session)
    params: Params = field(default_factory=Params)
    payloads: List = field(default_factory=[])

    def to_dict(self):
        return {
            "session": self.session.to_dict(),
            "params": self.params,
            "payloads": []
        }

### 针对不同类型的模型，Params.model_params不同，需定制化开发
### 阿里千问
@dataclass
class ModelParamsQWen(ModelParams):
### 选传参数
    use_tools: bool = False
    eos_token_id: Optional[int] = None
    bad_words_ids: list[list[int]] = field(default_factory=list)
    prefix_allowed_tokens_fn: str = "test_function"
    extend_model_params: Dict[str, Any] = field(default_factory=dict)
### 固定参数
    model_id: str = "001_int4" # qwen模型名称固定为 001_int4
    num_beams: int = 1 # 固定为 1
    num_beam_groups: int = 1 # 固定为 1
    do_sample: bool = True # 设置为 True 才能设置 top_p
### dify模型配置
    top_p: float = 0.8
    top_k: int = 10
    temperature: float = 0.3
    repetition_penalty: float = 1.0
    diversity_penalty: float = 0.0
    length_penalty: float = 0.0
    max_new_tokens: int = 512
    max_length: int = 512
    num_return_sequences: int = 1

    def to_dict(self):
        return {
            "use_tools": self.use_tools,
            "eos_token_id": self.eos_token_id,
            "bad_words_ids": self.bad_words_ids,
            "prefix_allowed_tokens_fn": self.prefix_allowed_tokens_fn,
            "extend_model_params": self.extend_model_params,
            "model_id": self.model_id,
            "num_beams": self.num_beams,
            "num_beam_groups": self.num_beam_groups,
            "do_sample": self.do_sample,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "temperature": self.temperature,
            "repetition_penalty": self.repetition_penalty,
            "diversity_penalty": self.diversity_penalty,
            "length_penalty": self.length_penalty,
            "max_new_tokens": self.max_new_tokens,
            "max_length": self.max_length,
            "num_return_sequences": self.num_return_sequences,
        }

### Deepseek
@dataclass
class ModelParamsDeepSeek(ModelParams):
### 必传参数4
    model_id: str
    model_secret: str
    timeout: int = 600
### dify模型配置
    top_p: float = 0.96
    top_k: int = 20
    temperature: float = 0.6
    repetition_penalty: float = 1
    max_tokens: int = 10240

    def to_dict(self):
        return {
            "model_id": self.model_id,
            "model_secret": self.model_secret,
            "timeout": self.timeout,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "temperature": self.temperature,
            "repetition_penalty": self.repetition_penalty,
            "max_tokens": self.max_tokens,
        }



# 请求头
def get_request_headers(app_id: str,
                        secret_id: str,
                        secret_key: str,
                        ability_id: str,
                        action: str = 'default',
                        scene_id: str = 'default',
                        http_method: str = 'POST'):
    random_str = str(random.randint(1, 99999999999))
    request_headers = {
        'X-IFAC-Ability': ability_id,
        'X-IFAC-Action': action,
        'X-IFAC-App-Id': app_id,
        'X-IFAC-Nonce': random_str,
        'X-IFAC-Scene-Id': scene_id,
        'X-IFAC-Timestamp': str(int(time.mktime(datetime.datetime.now().timetuple()))),
    }
    request_headers['Authorization'] = get_authorization_header(http_method, request_headers, secret_id, secret_key)
    return request_headers


# 请求头加密字段
def get_authorization_header(http_method: str,
                             canonicalized_headers: dict,
                             secret_id: str,
                             secret_key: str):
    signature_method = 'HmacSHA256'
    canonicalized_headers_str = ''
    for key in sorted(canonicalized_headers.keys()):
        canonicalized_headers_str += key.lower().strip() + '=' + str(canonicalized_headers[key]).strip() + '\n'
    canonicalized_resource_str = '/'
    string_to_sign = http_method.upper() + '\n' + canonicalized_headers_str + canonicalized_resource_str
    signature = base64.b64encode(
        hmac.new(key=secret_key.encode('utf-8'), msg=string_to_sign.encode('utf-8'), digestmod=sha256).digest()
    )
    signature = str(signature, 'utf-8')
    authorization = str(signature_method) + ':' + str(secret_id) + ':' + str(signature)
    return authorization


### 请求体
def get_request_body(
        prompt_messages:list[Message],
        credentials: dict,
        model_parameters: dict,
        is_init: bool = False,
    ) -> PiccRequest:
    # model = PiccModel.value_of(credentials["model_base"].lower())
    if is_init:
        model_parameters = {}
    print(f"prompt_messages:{prompt_messages}")
    messages = Messages(
            text=prompt_messages
        )
    print(f"model_parameters:{model_parameters}")
    params = {
        "prompt": {
            "header": {
                "traceId": ""
            },
            "parameter": {
                "compatible": True,
                "chat": {
                    "temperature": model_parameters.get("temperature",0.1)
                },
                "extension": {
                    "model": "Qwen2.5-72B-Instruct-GPTQ-Int4",
                    "presence_penalty": model_parameters.get("presence_penalty",1.03),
                    "frequency_penalty": model_parameters.get("frequency_penalty",1.0),
                    "seed": 0,
                    "max_tokens": model_parameters.get("max_tokens",2048),
                    "top_p": model_parameters.get("top_p",0.95),
                    "stream": True
                }
            },
            "payload": {
                "message": messages.to_dict()
            }
        }
    }
    print(f"params:{params}")
    session = Session(
        status=SessionStatus.ONCE
    )
    request_body = {
    "session": session.to_dict(),
    "params": params,
    "payloads": []
    }
    request_body_copy = copy.deepcopy(request_body)
    return request_body_copy
